#!/usr/bin/python

#Usage: pfam-cf-extra.py <start directory> <kazusa file>

import sys
import string
import os

if __name__ == "__main__":
	
	#Read the input parameters
	start_dir = sys.argv[1]
	kaz_file = sys.argv[2]

	#Open the kazusa file
	kaz = open(kaz_file,"r")
	
	#A list of organisms found to prevent duplicates
	#found_orgs = list()

	for dirname, dirnames, filenames in os.walk(start_dir):

		#For each pfam-directory look at the fasta file
		for subdirname in dirnames:
			fam_name = subdirname.replace("/","")
			org_id = ""
			print("Extracting CF for " + fam_name)
			found_orgs = list()
	
			#Open each fasta file
			fa_file = open(start_dir + "/" + subdirname + "/" + fam_name + ".fa","r")
			
			#Open an output file to write the codon frequences to
			outfile = open(start_dir + "/" + subdirname + "/" + fam_name + ".cf", "w")
	
			line = fa_file.readline()
			while line:
				line = line.rstrip()
				#Extract org_id from fasta header
				if line.startswith(">"):
					s_split = line.split(":")
					org_id = s_split[1].rstrip()
					if org_id not in found_orgs:
						found_orgs.append(org_id)
						
						#Search the kaz file for the org_id
						kaz.seek(0)
						k_line = kaz.readline()
						while k_line:
							k_split = k_line.split(":")
							if ":" in k_line and k_split[0] == org_id:
								k_line = kaz.readline()
								k_line = k_line.rstrip()
								outfile.write(line + "\n")
								outfile.write(k_line + "\n")
								break
	
							#Read the next line from kaz	
							k_line = kaz.readline()
					
				line = fa_file.readline()
			
			outfile.close()
			fa_file.close()		 
										
	kaz.close()

